"""
Abstract interfaces for model components in DIL_Frame.

This module defines the core interfaces that all model components must implement,
ensuring consistent API across different implementations.
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Optional

import torch
from torch import nn

from models.config import BackboneConfig


class BackboneBase(nn.Module, ABC):
    """Abstract base class for all backbone implementations.

    This class defines the core interface that all backbone networks must implement,
    ensuring consistent API across different implementations.
    """

    @abstractmethod
    def __init__(self, config: BackboneConfig) -> None:
        """Initialize the backbone with configuration.

        Args:
            config: Configuration object containing backbone parameters
        """
        super().__init__()
        self.config = config

    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the backbone.

        Args:
            x: Input tensor of shape [B, C, H, W]

        Returns:
            Feature tensor from the backbone
        """
        pass

    @abstractmethod
    def extract_features(
        self, x: torch.Tensor, layer_idx: Optional[List[int]] = None
    ) -> Dict[str, torch.Tensor]:
        """Extract features from specified layers of the backbone.

        Args:
            x: Input tensor of shape [B, C, H, W]
            layer_idx: Indices of layers to extract features from.
                       If None, extract features from default layers.

        Returns:
            Dictionary mapping layer names to feature tensors
        """
        pass

    def get_output_dim(self) -> int:
        """Get the output dimension of the backbone.

        Returns:
            Dimension of the output features
        """
        return self.config.output_dim


class AdapterBase(nn.Module, ABC):
    """Abstract base class for all parameter-efficient adapters.

    This class defines the interface for adapter modules used for
    parameter-efficient fine-tuning.
    """

    @abstractmethod
    def __init__(
        self, input_dim: int, bottleneck_dim: int, dropout: float = 0.0
    ) -> None:
        """Initialize the adapter.

        Args:
            input_dim: Dimension of input features
            bottleneck_dim: Dimension of bottleneck layer
            dropout: Dropout probability
        """
        super().__init__()
        self.input_dim = input_dim
        self.bottleneck_dim = bottleneck_dim
        self.dropout = dropout

    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the adapter.

        Args:
            x: Input tensor

        Returns:
            Adapted output tensor
        """
        pass

    def train(self, mode: bool = True) -> "AdapterBase":
        """Set training mode for adapter.

        Args:
            mode: Whether to set training mode (True) or evaluation mode (False)

        Returns:
            Self for chaining
        """
        super().train(mode)
        return self

    def get_param_count(self) -> int:
        """Get the number of parameters in the adapter.

        Returns:
            Number of trainable parameters
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)
